"""
优化版本的主程序
支持参数优化、WebSocket实时数据、自适应策略
"""
import os
import sys
from datetime import datetime, timedelta
from typing import List, Dict

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from src.utils.config import Config
from src.utils.logger import get_logger
from src.okx.client import OKXClient
from src.data.kline import KlineDataFetcher
from src.data.processor import KlineProcessor
from src.data.realtime_ws import RealtimeDataManager
from src.strategy.optimized_scalping import OptimizedScalpingStrategy
from src.strategy.parameter_optimizer import StrategyParams, ParameterOptimizer
from src.trading.executor import TradeExecutor
from src.trading.risk import RiskManager, RiskControl
from src.backtesting.optimized_engine import OptimizedBacktestEngine
from src.costs.calculator import CostCalculator

logger = get_logger(__name__)


class OptimizedScalpingBot:
    """优化版本的头皮交易机器人"""

    def __init__(self, config_path: str = None):
        """初始化机器人"""
        # 加载配置
        self.config = Config.load(config_path)

        # 初始化API客户端
        self.client = OKXClient(
            api_key=self.config.get('api.api_key', ''),
            secret_key=self.config.get('api.secret_key', ''),
            passphrase=self.config.get('api.passphrase', ''),
            environment=self.config.get('api.environment', 'sandbox'),
            simulate=self.config.get('simulate.enabled', False)
        )

        # 初始化数据处理
        self.fetcher = KlineDataFetcher(self.client)
        self.processor = KlineProcessor()

        # 初始化WebSocket连接
        self.realtime_manager = RealtimeDataManager(
            api_key=self.config.get('api.api_key', ''),
            secret_key=self.config.get('api.secret_key', ''),
            passphrase=self.config.get('api.passphrase', ''),
            use_sandbox=self.config.get('api.environment', 'sandbox') == 'sandbox'
        )

        # 添加数据回调
        self.realtime_manager.add_kline_callback(self._on_kline_update)
        self.realtime_manager.add_ticker_callback(self._on_ticker_update)
        self.realtime_manager.add_error_callback(self._on_ws_error)

        # 策略参数
        self.strategy_params = StrategyParams()
        self.strategy = OptimizedScalpingStrategy(self.config.get('strategy', {}), self.strategy_params)

        # 初始化交易执行
        self.executor = TradeExecutor(
            self.client,
            trade_type=self.config.get('trading.trade_type', 'spot'),
            leverage=self.config.get('trading.leverage', 1.0),
            cost_calculator=CostCalculator(
                maker_fee=self.config.get('costs.maker_fee', 0.02),
                taker_fee=self.config.get('costs.taker_fee', 0.05),
                funding_rate=self.config.get('costs.funding_rate', 0.03)
            )
        )

        # 初始化风险管理
        risk_control = RiskControl(
            stop_loss_pct=self.config.get('risk.stop_loss', 1.0),
            take_profit_pct=self.config.get('risk.take_profit', 2.0),
            max_positions=self.config.get('risk.max_positions', 3),
            max_daily_loss_pct=self.config.get('risk.max_daily_loss', 5.0),
            max_drawdown_pct=self.config.get('risk.max_drawdown', 10.0),
            max_position_bars=self.config.get('trading.max_position_bars', 20),
        )
        self.risk_manager = RiskManager(risk_control, self.config.get('trading.position_size', 100))

        # 性能统计
        self.performance_stats = {
            'signals_generated': 0,
            'trades_executed': 0,
            'last_signal_time': None,
            'websocket_messages_received': 0,
        }

        logger.info("优化版本的机器人初始化完成")

    def _on_kline_update(self, symbol: str, timeframe: str, kline_data: Dict):
        """K线数据更新回调"""
        self.performance_stats['websocket_messages_received'] += 1

        # 只在K线确认时处理 (confirm == '1')
        if kline_data.get('confirm') != '1':
            return

        logger.debug(f"K线更新: {symbol} {timeframe} @ {kline_data['close']}")

        # 根据K线更新执行策略
        self._check_trading_signal(symbol)

    def _on_ticker_update(self, ticker_data: Dict):
        """行情数据更新回调"""
        # 可以在这里添加实时价格更新的逻辑
        pass

    def _on_ws_error(self, error):
        """WebSocket错误回调"""
        logger.error(f"WebSocket错误: {error}")

    def _check_trading_signal(self, symbol: str):
        """检查交易信号"""
        # 获取各时间框架的K线数据
        kline_data = {}
        for timeframe in self.config.get('strategy.timeframes', [5, 15, 60, 240]):
            klines = self.realtime_manager.get_klines(symbol, f"{timeframe}m" if timeframe < 60 else f"{timeframe//60}H")
            if klines:
                df = self.processor.add_klines(symbol, timeframe, klines)
                df = self.processor.add_technical_indicators(df, self.strategy_params.to_dict())
                kline_data[timeframe] = df

        if not kline_data:
            return

        # 生成信号
        signal, confidence = self.strategy.generate_signal(kline_data)

        if signal in ["buy", "sell"]:
            self.performance_stats['signals_generated'] += 1
            logger.info(f"生成信号: {signal.upper()} {symbol} (信心度: {confidence:.2f})")

            # 检查现仓位
            position = self.executor.position_manager.get_active_position(symbol)

            if position and signal != "hold":
                # 平仓
                self.executor.close_position(symbol, reason="signal")
            elif not position and signal in ["buy", "sell"]:
                # 开仓
                position_size = self.config.get('trading.position_size', 100)
                current_price = self.realtime_manager.get_ticker(symbol)
                if current_price:
                    price = current_price.get('last_price', 0)
                    quantity = position_size / price if price > 0 else 0

                    if signal == "buy":
                        success, trade = self.executor.execute_buy(symbol, quantity, order_type="market")
                    else:
                        success, trade = self.executor.execute_short(symbol, quantity, order_type="market")

                    if success:
                        self.performance_stats['trades_executed'] += 1
                        self.performance_stats['last_signal_time'] = datetime.now()

    def optimize_parameters(self):
        """参数优化"""
        logger.info("="*80)
        logger.info("开始参数优化")
        logger.info("="*80)

        # 获取回测配置
        backtest_config = self.config.get('backtest', {})
        start_date = datetime.strptime(backtest_config['start_date'], '%Y-%m-%d')
        end_date = datetime.strptime(backtest_config['end_date'], '%Y-%m-%d')
        symbols = backtest_config.get('symbols', ['BTC-USDT'])
        initial_balance = backtest_config.get('initial_balance', 10000)

        # 获取历史K线数据
        logger.info(f"获取历史K线数据: {start_date} - {end_date}")
        klines_data = {}
        for symbol in symbols:
            logger.info(f"获取 {symbol} 的历史数据...")
            klines_data[symbol] = {}

            for timeframe in self.config.get('strategy.timeframes', [5, 15, 60, 240]):
                try:
                    klines = self.fetcher.fetch_historical_klines(
                        symbol, timeframe, start_date, end_date
                    )
                    if klines:
                        df = self._convert_klines_to_df(klines)
                        df = self.processor.add_technical_indicators(df, {})
                        klines_data[symbol][timeframe] = df
                        logger.info(f"  {timeframe}分钟: {len(df)}条K线")
                except Exception as e:
                    logger.error(f"获取 {symbol} {timeframe}分钟K线失败: {e}")

        # 生成参数组合
        optimizer = ParameterOptimizer()
        params_list = optimizer.generate_parameter_combinations(
            ema_fast_range=(5, 15, 2),      # EMA快线: 5-13, 步长2
            ema_slow_range=(15, 35, 3),     # EMA慢线: 15-32, 步长3
            rsi_period_range=(10, 20, 2),   # RSI周期: 10-18, 步长2
            stop_loss_range=(0.005, 0.015, 0.002)  # 止损: 0.5%-1.5%, 步长0.2%
        )

        # 运行回测
        logger.info(f"开始网格搜索回测 ({len(params_list)} 个参数组合)...")
        backtest_engine = OptimizedBacktestEngine(self.config, initial_balance)

        for params in params_list:
            backtest_engine.current_params = params

        optimizer, results = backtest_engine.run_backtest_grid(
            symbols, klines_data, params_list, start_date, end_date
        )

        # 保存结果
        optimizer.save_results("optimization_results.json")

        # 返回最优参数
        return optimizer.best_params

    def run_backtest(self):
        """运行回测"""
        logger.info("="*80)
        logger.info("启动回测模式")
        logger.info("="*80)

        # 获取回测配置
        backtest_config = self.config.get('backtest', {})
        start_date = datetime.strptime(backtest_config['start_date'], '%Y-%m-%d')
        end_date = datetime.strptime(backtest_config['end_date'], '%Y-%m-%d')
        symbols = backtest_config.get('symbols', ['BTC-USDT'])
        initial_balance = backtest_config.get('initial_balance', 10000)

        # 获取历史K线数据
        klines_data = {}
        for symbol in symbols:
            logger.info(f"获取 {symbol} 的历史数据...")
            klines_data[symbol] = {}

            for timeframe in self.config.get('strategy.timeframes', [5, 15, 60, 240]):
                try:
                    klines = self.fetcher.fetch_historical_klines(
                        symbol, timeframe, start_date, end_date
                    )
                    if klines:
                        df = self._convert_klines_to_df(klines)
                        df = self.processor.add_technical_indicators(df, self.strategy_params.to_dict())
                        klines_data[symbol][timeframe] = df
                except Exception as e:
                    logger.error(f"获取K线失败: {e}")

        # 运行回测
        backtest_engine = OptimizedBacktestEngine(self.config, initial_balance)
        report = backtest_engine.run_backtest_single(
            symbols, klines_data, self.strategy_params, start_date, end_date
        )

        logger.info(f"\n回测完成: {report}")
        return report

    def run_realtime(self):
        """运行实时交易 (使用WebSocket)"""
        logger.info("="*80)
        logger.info("启动实时交易模式 (WebSocket)")
        logger.info("="*80)

        symbols = self.config.get('trading.symbols', ['BTC-USDT', 'ETH-USDT'])
        timeframes = self.config.get('strategy.timeframes', [5, 15, 60, 240])

        # 转换时间框架格式
        ws_timeframes = []
        for tf in timeframes:
            if tf < 60:
                ws_timeframes.append(f"{tf}m")
            elif tf == 60:
                ws_timeframes.append("1H")
            elif tf == 240:
                ws_timeframes.append("4H")

        try:
            # 初始化WebSocket连接
            logger.info("初始化WebSocket连接...")
            self.realtime_manager.init_public_data(symbols, ws_timeframes)
            self.realtime_manager.wait_ready(timeout=30)

            logger.info("WebSocket连接就绪，开始实时交易...")

            # 主循环
            while True:
                # 定期输出性能统计
                logger.info(
                    f"性能统计: "
                    f"WebSocket消息={self.performance_stats['websocket_messages_received']}, "
                    f"生成信号={self.performance_stats['signals_generated']}, "
                    f"执行交易={self.performance_stats['trades_executed']}"
                )

                import time
                time.sleep(60)

        except KeyboardInterrupt:
            logger.info("用户中断交易")
        except Exception as e:
            logger.error(f"交易异常: {e}")
        finally:
            self.realtime_manager.disconnect()

    def _convert_klines_to_df(self, klines) -> 'pd.DataFrame':
        """将K线对象列表转换为DataFrame"""
        import pandas as pd

        data = {
            'timestamp': [k.timestamp for k in klines],
            'open': [k.open for k in klines],
            'high': [k.high for k in klines],
            'low': [k.low for k in klines],
            'close': [k.close for k in klines],
            'volume': [k.volume for k in klines],
        }

        df = pd.DataFrame(data)
        df = df.sort_values('timestamp').reset_index(drop=True)
        return df


def main():
    """主入口"""
    import argparse

    parser = argparse.ArgumentParser(description='优化版本的头皮交易机器人')
    parser.add_argument('--mode', choices=['backtest', 'optimize', 'realtime'],
                       default='backtest', help='运行模式')
    parser.add_argument('--config', default=None, help='配置文件路径')

    args = parser.parse_args()

    # 初始化机器人
    bot = OptimizedScalpingBot(args.config)

    # 根据模式运行
    if args.mode == 'backtest':
        bot.run_backtest()

    elif args.mode == 'optimize':
        best_params = bot.optimize_parameters()
        if best_params:
            logger.info(f"\n最优参数:")
            logger.info(f"  EMA: {best_params.ema_fast_period}/{best_params.ema_slow_period}")
            logger.info(f"  RSI: {best_params.rsi_period}")
            logger.info(f"  止损: {best_params.stop_loss_pct:.2%}")
            logger.info(f"  止盈: {best_params.take_profit_pct:.2%}")

    elif args.mode == 'realtime':
        bot.run_realtime()


if __name__ == '__main__':
    main()
